import numpy as np
import geatpy as ea


class myProblem(ea.Problem):
    def __init__(self):
        name = 'myProblem'
        M = 2
        maxormins = [1] * M
        Dim = 2
        varTypes = [0] * Dim
        lb = [-5] * Dim
        ub = [5] * Dim
        lbin = [1] * Dim
        ubin = [1] * Dim
        ea.Problem.__init__(self, name, M, maxormins, Dim, varTypes, lb, ub, lbin, ubin)

    def aimFunc(self, pop):
        x1 = pop.Phen[:, [0]]
        x2 = pop.Phen[:, [1]]
        pop.ObjV = np.zeros((pop.Phen.shape[0], self.M))
        pop.ObjV[:, [0]] = x1**4 - 10*x1**2 + x1*x2 + x2**4 - x1**2*x2**2
        pop.ObjV[:, [1]] = x2**4 - x1**2*x2**2 + x1**4 + x1*x2


if __name__ == '__main__':
    problem = myProblem()

    algorithm = ea.moea_NSGA2_templet(problem, ea.Population(Encoding='RI', NIND=1000),
                                      MAXGEN=100,
                                      logTras=0)

    res = ea.optimize(algorithm, verbose=False, drawing=1, outputMsg=True, drawLog=False, saveFlag=False, dirName='result')
    print(res)

